# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import json
import logging
import traceback

from absl import app, flags
from tornado.ioloop import IOLoop
from pymongo import MongoClient

from coin.base.timestamp import get_timestamp
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.util.queue.config import KafkaConfig
from coin.util.queue.reader import AsyncKafkaReader, kafka
from coin.experimental.xguo.fruit.util.util import OhlcCalculator, PeriodicAccumulator

FLAGS = flags.FLAGS


class KafkaMonitor(object):
  def __init__(self, *, kafka_config: KafkaConfig):
    self._kafka_config = kafka_config
    self._ioloop = None
    self._logger = logging.getLogger(__name__)
    self._ohlc_calculator = OhlcCalculator(window_sec=600)
    self._accumulator = PeriodicAccumulator(window_sec=600)
    self._db = MongoClient().candle_20190101
    """
    self._db = MongoClient(
        'mongodb://candle_importer:GAX#GfvX4G$!f%40vf'
        '@coin-mongodb.corp.prestolabs.io/candle_20190101').candle_20190101
    """

  async def _loop(self):
    topics = self.query_topics()
    self._logger.info('Kafka Topics: \n%s', json.dumps(topics, indent=2))
    offset_nano_sec = datetime.timedelta(hours=0).total_seconds() * 1e9
    start_ts = get_timestamp() - int(offset_nano_sec)
    reader = AsyncKafkaReader(topics, kafka_config=self._kafka_config, timestamp_from=start_ts)
    await reader.open()
    try:
      while True:
        # noinspection PyBroadException
        try:
          records = await reader.read()
          for msg in records:
            timestamp = int(msg.timestamp * 1e6)
            topic = msg.topic
            pb = StrategyLog()
            pb.ParseFromString(msg.value)
            self.on_message(timestamp, topic, pb)
        except ConnectionError:
          self._logger.debug('Connection error!')
          continue
        except Exception:
          self._logger.error(traceback.format_exc())
          continue
    finally:
      await reader.close()
      self._ioloop.stop()

  def query_topics(self):
    servers = self._kafka_config.kafka_servers
    consumer = kafka.KafkaConsumer(bootstrap_servers=servers)
    topics = consumer.topics()
    topics = [topic for topic in topics if topic.startswith('STRAT') and 'trimm' in topic]
    return topics

  def on_message(self, timestamp, topic, data):
    self._logger.debug('%s, %s, data', timestamp, topic, data)

    ohlc = self._ohlc_calculator.emit(timestamp)
    if ohlc is not None:
      fill = self._accumulator.emit(timestamp)
      if fill is None:
        volume = 0.
      else:
        volume = fill.value

      msg = {
          'revision':
              '20190610',
          'timestamp':
              datetime.datetime.fromtimestamp(timestamp / 1e9),
          'payload': [{
              'open': ohlc.open,
              'close': ohlc.close,
              'high': ohlc.high,
              'low': ohlc.low,
              'volume': volume,
              'volume_agg': 0,
              'bases': {},
              'value': 1.0,
              'currency': 'BTC',
          }]
      }
      self._db['unhedge_trimm_okex_xrp'].insert_one(msg)

    if data.type == StrategyLog.PNL_BALANCE:
      self._ohlc_calculator.update(timestamp, data.pnl_balance.total)
    elif data.type == StrategyLog.OG_LOG:
      og_log = data.og_log
      event = og_log.event
      if event.type not in (OrderEvent.ORDER_FILLED,):
        return
      self._accumulator.update(timestamp, event.fill_qty)

      self._logger.info(
          'FILL %s\t%f\t%f\t%f',
          event.symbol,
          event.fill_price,
          event.fill_qty,
          event.order_side,
      )

  def start(self):
    self._ioloop = IOLoop.current()
    self._ioloop.add_callback(self._loop)
    self._ioloop.start()


def main(_):
  config = KafkaConfig.from_kafka_servers('coin-kafka.corp.prestolabs.io:9092')
  monitor = KafkaMonitor(kafka_config=config)
  monitor.start()


if __name__ == '__main__':
  flags.DEFINE_boolean(
      'use_local_mongo',
      False,
      'Use local mongo or not',
  )
  app.run(main)
